#include "my_disk_hotspot.h"

#include "comm/my_log.h"
#include "comm/my_mem.h"

typedef struct {
    int                 idx;
    uint32_t            digest32;
    int64_t             offset;
    int64_t             size;
    my_queue_t          queue;
    my_rbtree_node_t    tree_node;
} my_disk_hotspot_item_t;

static inline my_rbtree_node_t *
my_disk_hotspot_lookup_node(my_rbtree_t *tree, int64_t offset) {
    my_rbtree_node_t *node = tree->root;
    my_rbtree_node_t *sentinel = tree->sentinel;

    while (node != sentinel) {
        if (offset > node->key) {
            node = node->right;
            continue;
        }

        if (offset < node->key) {
            node = node->left;
            continue;
        }

        return node;
    }

    return NULL; // not found
}

static my_bool_t
my_disk_hotspot_list_is_full(my_disk_hotspot_t *hotspot, int list_idx) {
    if (*(hotspot->lists_n + list_idx) >= hotspot->max_item) {
        return MY_TRUE;
    }

    if (*(hotspot->lists_bytes + list_idx) >= hotspot->max_bytes) {
        return MY_TRUE;
    }

    return MY_FALSE;
}

static void
my_disk_hotspot_insert_record(my_disk_hotspot_t *hotspot, const my_disk_hotspot_record_t * record) {
    my_disk_hotspot_item_t *item = NULL;

    if (my_disk_hotspot_list_is_full(hotspot, 0)) {
        my_queue_t *back_node = my_queue_last((hotspot->lists + 0));
        my_disk_hotspot_item_t *item =
            my_queue_data(back_node, my_disk_hotspot_item_t, queue);
        *(hotspot->lists_n + 0) -= 1;
        *(hotspot->lists_bytes + 0) -= item->size;
        my_queue_remove(&item->queue);
        my_rbtree_delete(&hotspot->tree, &item->tree_node);
    }

    if (!item) {
        item = my_malloc(sizeof(my_disk_hotspot_item_t));
    }

    item->idx = 0;
    item->digest32 = record->digest32;
    item->offset = record->offset;
    item->size = record->size;
    item->tree_node.key = record->offset;

    *(hotspot->lists_n + 0) += 1;
    *(hotspot->lists_bytes + 0) += item->size;
    my_rbtree_insert(&hotspot->tree, &item->tree_node);
    my_queue_insert_head(hotspot->lists + 0, &item->queue);
}

static void
my_disk_hotspot_increase_priority_helper(my_disk_hotspot_t *hotspot,
    const my_disk_hotspot_record_t * record) {

    my_rbtree_node_t *node;
    node = my_disk_hotspot_lookup_node(&hotspot->tree, record->offset);
    if (!node) {
        my_disk_hotspot_insert_record(hotspot, record);
        return;
    }

    my_disk_hotspot_item_t *item = 
        (my_disk_hotspot_item_t *)((char *)node - offsetof(my_disk_hotspot_item_t, tree_node));

    // support record digest and size update
    *(hotspot->lists_bytes + item->idx) -= item->size;
    *(hotspot->lists_bytes + item->idx) += record->size;
    item->size = record->size;
    item->digest32 = record->digest32;

    // already on final list? move to header.
    if (item->idx == hotspot->nlists - 1) {
        my_queue_remove(&item->queue);
        my_queue_insert_head((hotspot->lists + item->idx), &item->queue);
        return;
    }

    *(hotspot->lists_bytes + item->idx) -= item->size;
    *(hotspot->lists_n + item->idx) -= 1;
    my_queue_remove(&item->queue);
    item->idx += 1;

    if (!my_disk_hotspot_list_is_full(hotspot, item->idx)) {
        *(hotspot->lists_bytes + item->idx) += item->size;
        *(hotspot->lists_n + item->idx) += 1;
        my_queue_insert_head((hotspot->lists + item->idx), &item->queue);
        return;
    }

    // current list is full, swap
    my_queue_t *back_node = my_queue_last((hotspot->lists + item->idx));
    my_disk_hotspot_item_t *back_item =
        my_queue_data(back_node, my_disk_hotspot_item_t, queue);
    my_queue_remove(back_node);
    *(hotspot->lists_bytes + back_item->idx) -= back_item->size;
    *(hotspot->lists_n + back_item->idx) -= 1;
    back_item->idx -= 1;
    *(hotspot->lists_bytes + back_item->idx) += back_item->size;
    *(hotspot->lists_n + back_item->idx) += 1;
    my_queue_insert_head((hotspot->lists + back_item->idx), &back_item->queue);

    *(hotspot->lists_bytes + item->idx) += item->size;
    *(hotspot->lists_n + item->idx) += 1;
    my_queue_insert_head((hotspot->lists + item->idx), &item->queue);
}

void
my_disk_hotspot_increase_priority(my_disk_hotspot_t *hotspot,
    const my_disk_hotspot_record_t * record) {

    my_spin_lock_lock(&hotspot->lock);
    my_disk_hotspot_increase_priority_helper(hotspot, record);
    my_spin_lock_unlock(&hotspot->lock);
}

// return the first node greater than or equal to offset
static my_rbtree_node_t *
my_disk_hotspot_first_largest(my_rbtree_node_t *node, my_rbtree_node_t *sentinel,
    int64_t window_off) {

    my_rbtree_node_t *save = NULL;
    while (node != sentinel) {
        if (node->key == window_off) {
            return node;
        }

        if (node->key < window_off) {
            node = node->right;
            continue;
        }

        save = node;
        node = node->left;
    }

    return save;
}

static my_bool_t
my_disk_hotspot_first_cover_helper(my_disk_hotspot_t *hotspot, int64_t window_off,
    int64_t window_size, my_disk_hotspot_record_t *record) {

    my_rbtree_node_t *node = my_disk_hotspot_first_largest(hotspot->tree.root,
        hotspot->tree.sentinel, window_off);

    if (!node) {
        return MY_FALSE;
    }

    my_disk_hotspot_item_t *item = 
        (my_disk_hotspot_item_t *)((char *)node - offsetof(my_disk_hotspot_item_t, tree_node));

    if (item->offset > window_off + window_size) {
        return MY_FALSE;
    }

    record->digest32 = item->digest32;
    record->offset = item->offset;
    record->size = item->size;
    return MY_TRUE;
}

my_bool_t
my_disk_hotspot_first_cover(my_disk_hotspot_t *hotspot,int64_t window_off,
    int64_t window_size, my_disk_hotspot_record_t *record) {

    my_bool_t res;
    my_spin_lock_lock(&hotspot->lock);
    res = my_disk_hotspot_first_cover_helper(hotspot, window_off, window_size, record);
    my_spin_lock_unlock(&hotspot->lock);
    return res;
}

void
my_disk_hotspot_open(my_disk_hotspot_t *hotspot, int nlists, int64_t max_item,
    int64_t max_bytes) {

    if (nlists <= 0) {
        MY_LOG_FATAL0("disk_hotspot", "nlist must be greater than 0");
    }

    bzero(hotspot, sizeof(my_disk_hotspot_t));
    my_spin_lock_init(&hotspot->lock);

    hotspot->max_item = max_item / nlists;
    hotspot->max_bytes = max_bytes / nlists;
    if (hotspot->max_item <= 0 || hotspot->max_bytes <= 0) {
        MY_LOG_FATAL0("disk_hotspot", "max_item and max_bytes must be greater than 0");
    }

    my_queue_t *lists = my_malloc(sizeof(my_queue_t) * nlists);
    for (int i = 0; i < nlists; i++) {
        my_queue_init(lists + i);
    }

    int64_t *lists_bytes = my_malloc(sizeof(int64_t) * nlists);
    bzero(lists_bytes, sizeof(int64_t) * nlists);

    int64_t *lists_n = my_malloc(sizeof(int64_t) * nlists);
    bzero(lists_n, sizeof(int64_t) * nlists);

    hotspot->nlists = nlists;
    hotspot->lists = lists;
    hotspot->lists_bytes = lists_bytes;
    hotspot->lists_n = lists_n;

    my_rbtree_init(&hotspot->tree, &hotspot->sentinel, my_rbtree_insert_value);
}

void
my_disk_hotspot_close(my_disk_hotspot_t *hotspot) {
    my_queue_t  *q;
    my_disk_hotspot_item_t *item; 

    my_spin_lock_lock(&hotspot->lock);
    for (int i = 0; i < hotspot->nlists; i++) {
        while (!my_queue_empty(hotspot->lists + i)) {
            q = my_queue_last(hotspot->lists+i);
            item = my_queue_data(q, my_disk_hotspot_item_t, queue);
            my_queue_remove(q);
            my_rbtree_delete(&hotspot->tree, &item->tree_node);
            my_free(item);
        }
    }
    my_free(hotspot->lists_n);
    my_free(hotspot->lists_bytes);
    my_free(hotspot->lists);
    my_spin_lock_unlock(&hotspot->lock);
}
